import uuid
from typing import Optional

from fastapi import APIRouter, Body
from pydantic import BaseModel, Field

from app.core.chains import get_rag_chain
from app.core.memory.memory_manager import get_session_history

router = APIRouter()


class KBChatRequest(BaseModel):
    query: str = Field(..., description="User's question for the knowledge base.")
    conversation_id: Optional[str] = Field(None,
                                           description="A unique ID for the conversation. If not provided, a new one will be generated.")


class KBChatResponse(BaseModel):
    answer: str
    conversation_id: str


@router.post("/", response_model=KBChatResponse)
def knowledge_base_chat(request: KBChatRequest = Body(...)):
    """
    Handles a stateful chat request against the knowledge base.
    Uses conversation_id to maintain context.
    """
    conv_id = request.conversation_id or str(uuid.uuid4())
    print(f"--- KB Chat endpoint called for conversation_id: {conv_id} ---")

    chat_history_backend = get_session_history(conv_id)
    rag_chain = get_rag_chain()

    result = rag_chain.invoke({
        "input": request.query,
        "chat_history": chat_history_backend.messages
    })

    answer = result.get("answer", "I'm sorry, I couldn't find an answer.")

    chat_history_backend.add_user_message(request.query)
    chat_history_backend.add_ai_message(answer)

    return KBChatResponse(
        answer=answer,
        conversation_id=conv_id
    )



